import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras.models import Sequential,Model
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.layers import Conv2D,BatchNormalization,ReLU,MaxPool2D,Flatten,Dense
from tensorflow.keras.losses import categorical_crossentropy
from tensorflow.keras.utils import to_categorical

import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

(_,_),(x_test,y_test) = cifar10.load_data()

x_test = x_test/255
y_test = to_categorical(y_test,10)

x_train,x_test,y_train,y_test = train_test_split(x_test,y_test,train_size=0.8)
x_val,x_test,y_val,y_test = train_test_split(x_test,y_test,train_size=0.5)
print(x_train.shape)
print(x_test.shape)
print(x_val.shape)

class ConvCell(Model):
    def __init__(self,filters,kernel_size):
        super(ConvCell, self).__init__()
        self.conv = Sequential([
            Conv2D(filters,kernel_size,strides=1,padding='same'),
            BatchNormalization(),
            ReLU()
        ])
    def call(self, inputs, training=None, mask=None):
        return self.conv(inputs)

class VGG16(Model):
    cfg = [(64, 2), (128, 2), (256, 3), (512, 3), (512, 3)]
    def __init__(self):
        super(VGG16, self).__init__()
        self.conv_layers = self.make_layers()
        self.flatten = Flatten()
        self.fc1 = Dense(256,activation='relu')
        self.fc2 = Dense(64,activation='relu')
        self.fc3 = Dense(10,activation='softmax')

    def make_layers(self):
        layer = Sequential()
        for c in self.cfg:
            filters = c[0]
            for i in range(c[1]):
                layer.add(ConvCell(filters,3))
            layer.add(MaxPool2D())
        return layer
    def call(self, inputs, training=None, mask=None):
        x = self.conv_layers(inputs)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

if __name__ == '__main__':
    model = VGG16()
    model.build((None,32,32,3))
    model.summary()

    model.compile(optimizer='adam',loss=categorical_crossentropy,metrics=['accuracy'])
    history = model.fit(x_train,y_train,epochs=10,validation_data=(x_val,y_val))

    plt.plot(history.history['loss'], 'ro-', label='train')
    plt.plot(history.history['val_loss'], 'go--', label='validation')
    plt.legend()
    plt.title('loss')
    plt.show()

    plt.plot(history.history['accuracy'], 'ro-', label='train')
    plt.plot(history.history['val_accuracy'], 'go--', label='validation')
    plt.legend()
    plt.title('accuracy')
    plt.show()

    score = model.evaluate(x_test, y_test,verbose=0)
    print(f'测试集损失值:{score[0]:.3f}')
    print(f'测试集准确率:{score[1]:.3f}')

    plt.pie([score[1], 1 - score[1]], [0, 0.01], ['pred_true', 'pred_false'], autopct='%1.1f%%')
    plt.show()